//! Loop statement parsing (for, while)

use super::super::{parse_expression, Parser};
use super::blocks::parse_block_statement;
use crate::ast::stmt::{ForStmt, ForVariant, Statement, WhileStmt};
use crate::error::{CompilerError, Result};
use crate::lexer::token::Token;

/// Parse for statement - supports Go-style for loops
pub fn parse_for_statement(parser: &mut Parser) -> Result<ForStmt> {
    parser.expect(&Token::For)?;

    // Check if this is a traditional C-style for loop: for init; cond; inc {}
    // Look for semicolon to determine if this is a traditional for loop
    if let Some(Token::Semicolon) = parser.peek().map(|t| &t.token) {
        // This is a traditional for loop with no initializer
        parser.advance()?; // consume ';'

        // Parse condition (optional)
        let condition = if parser.check(&Token::Semicolon) {
            None
        } else {
            let expr = parse_expression(parser)?;
            if expr.is_none() {
                return Err(CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected expression in for loop condition",
                ));
            }
            Some(expr.unwrap())
        };

        // Expect semicolon after condition
        parser.expect(&Token::Semicolon)?;

        // Parse increment (optional)
        let increment = if parser.check(&Token::LeftBrace) {
            None
        } else {
            let stmt = super::parse_statement(parser)?;
            if stmt.is_none() {
                return Err(CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected statement in for loop increment",
                ));
            }
            Some(Box::new(stmt.unwrap()))
        };

        // Parse loop body
        let body = if parser.check(&Token::LeftBrace) {
            parse_block_statement(parser).map(|block| Statement::Block(block))?
        } else {
            super::parse_statement(parser)?.ok_or_else(|| {
                CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected statement after for loop",
                )
            })?
        };

        Ok(ForStmt {
            variant: ForVariant::Traditional {
                initializer: None,
                condition,
                increment,
            },
            body: Box::new(body),
            location: parser.current_location(),
        })
    }
    // Check if this starts with 'var' keyword
    else if let Some(Token::Var) = parser.peek().map(|t| &t.token) {
        // Look ahead to determine if this is range-based (for var v T in slice)
        // or traditional (for var i = 0; i < 10; i++)
        // Three possible patterns:
        // 1. for var i = 0; ...           (traditional, no type)
        // 2. for var i int32 = 0; ...     (traditional, with type)
        // 3. for var v int32 in slice ... (range-based)

        // We need to distinguish between patterns 2 and 3
        // Strategy: Look for 'in' keyword in the token stream before we hit ';' or '='
        // If we find 'in' first, it's range-based; otherwise it's traditional

        let mut lookahead_index = 1; // Start after 'var'
        let mut is_range_based = false;

        // Skip identifier (variable name)
        if matches!(
            parser.peek_ahead(lookahead_index).map(|t| &t.token),
            Some(Token::Identifier(_))
        ) {
            lookahead_index += 1;

            // Now check what comes next:
            // - If '=' → traditional (pattern 1)
            // - If type followed by '=' → traditional (pattern 2)
            // - If type followed by 'in' → range-based (pattern 3)

            if matches!(
                parser.peek_ahead(lookahead_index).map(|t| &t.token),
                Some(Token::Equal)
            ) {
                // Pattern 1: for var i = 0
                is_range_based = false;
            } else {
                // Might be a type, so skip it
                // Types can be: Identifier (for custom types), primitive types, or complex types
                // For simplicity, skip tokens until we find 'in' or '='
                while lookahead_index < 10 {
                    // Limit lookahead to prevent infinite loop
                    match parser.peek_ahead(lookahead_index).map(|t| &t.token) {
                        Some(Token::In) => {
                            is_range_based = true;
                            break;
                        }
                        Some(Token::Equal) | Some(Token::Semicolon) => {
                            is_range_based = false;
                            break;
                        }
                        None => break,
                        _ => lookahead_index += 1,
                    }
                }
            }
        }

        let is_traditional = !is_range_based;

        if !is_traditional {
            // This is a range-based for loop: for var v T in vec {}
            parser.advance()?; // consume 'var'

            // Parse variable name
            let variable =
                if let Some(Token::Identifier(var_name)) = parser.peek().map(|t| &t.token) {
                    let name = var_name.clone();
                    parser.advance()?;
                    name
                } else {
                    return Err(CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected variable name after 'var'",
                    ));
                };

            // Parse variable type (required for now)
            let variable_type =
                super::super::type_parser::parse_type(parser)?.ok_or_else(|| {
                    CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected type after variable name in for-in loop",
                    )
                })?;

            // Expect 'in' keyword
            parser.expect(&Token::In)?;

            // Parse iterable expression
            let iterable = parse_expression(parser)?;
            if iterable.is_none() {
                return Err(CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected iterable after 'in'",
                ));
            }

            // Parse loop body
            let body = if parser.check(&Token::LeftBrace) {
                parse_block_statement(parser).map(|block| Statement::Block(block))?
            } else {
                super::parse_statement(parser)?.ok_or_else(|| {
                    CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected statement after for loop",
                    )
                })?
            };

            Ok(ForStmt {
                variant: ForVariant::Range {
                    variable,
                    variable_type,
                    iterable: iterable.unwrap(),
                },
                body: Box::new(body),
                location: parser.current_location(),
            })
        } else {
            // This is a traditional for loop with initializer: for var i = 0; i < 10; i++
            let stmt = super::parse_statement(parser)?;
            if stmt.is_none() {
                return Err(CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected statement in for loop initializer",
                ));
            }
            let initializer = Some(Box::new(stmt.unwrap()));

            // Parse condition (optional)
            let condition = if parser.check(&Token::Semicolon) {
                None
            } else {
                let expr = parse_expression(parser)?;
                if expr.is_none() {
                    return Err(CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected expression in for loop condition",
                    ));
                }
                Some(expr.unwrap())
            };

            // Expect semicolon after condition
            parser.expect(&Token::Semicolon)?;

            // Parse increment (optional)
            let increment = if parser.check(&Token::LeftBrace) {
                None
            } else {
                let stmt = super::parse_statement(parser)?;
                if stmt.is_none() {
                    return Err(CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected statement in for loop increment",
                    ));
                }
                Some(Box::new(stmt.unwrap()))
            };

            // Parse loop body
            let body = if parser.check(&Token::LeftBrace) {
                parse_block_statement(parser).map(|block| Statement::Block(block))?
            } else {
                super::parse_statement(parser)?.ok_or_else(|| {
                    CompilerError::syntax(
                        parser.current_location().line,
                        parser.current_location().column,
                        "Expected statement after for loop",
                    )
                })?
            };

            Ok(ForStmt {
                variant: ForVariant::Traditional {
                    initializer,
                    condition,
                    increment,
                },
                body: Box::new(body),
                location: parser.current_location(),
            })
        }
    }
    // This is a condition-only for loop: for condition {}
    else {
        // Parse condition expression
        let condition = parse_expression(parser)?;
        if condition.is_none() {
            return Err(CompilerError::syntax(
                parser.current_location().line,
                parser.current_location().column,
                "Expected condition expression",
            ));
        }

        // Parse loop body
        let body = if parser.check(&Token::LeftBrace) {
            parse_block_statement(parser).map(|block| Statement::Block(block))?
        } else {
            super::parse_statement(parser)?.ok_or_else(|| {
                CompilerError::syntax(
                    parser.current_location().line,
                    parser.current_location().column,
                    "Expected statement after for loop",
                )
            })?
        };

        Ok(ForStmt {
            variant: ForVariant::Condition {
                condition: condition.unwrap(),
            },
            body: Box::new(body),
            location: parser.current_location(),
        })
    }
}

/// Parse while statement
pub fn parse_while_statement(parser: &mut Parser) -> Result<WhileStmt> {
    parser.expect(&Token::While)?;

    let condition = parse_expression(parser)?;
    if condition.is_none() {
        return Err(CompilerError::syntax(
            parser.current_location().line,
            parser.current_location().column,
            "Expected condition after 'while'",
        ));
    }

    let body = super::parse_statement(parser)?;
    if body.is_none() {
        return Err(CompilerError::syntax(
            parser.current_location().line,
            parser.current_location().column,
            "Expected statement after 'while' condition",
        ));
    }

    Ok(WhileStmt {
        condition: condition.unwrap(),
        body: Box::new(body.unwrap()),
        location: parser.current_location(),
    })
}
